# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import coin.exchange.base.kr_rest.private_client_base as privcb
import coin.exchange.huobi.kr_rest.native_private_client as nprivc
from coin.exchange.huobi.kr_rest.currency import HuobiCurrency
from coin.exchange.huobi.kr_rest.product import HuobiProduct
from coin.exchange.huobi.kr_rest.native_private_client import HuobiAccountType

import coin.proto.coin_order_enums_pb2 as coin_order
from coin.proto.coin_query_pb2 import (AccountBalance,
                                       CurrencyBalance,
                                       CurrencyTransferBundle,
                                       CurrencyTransferElement,
                                       ProductFillElement,
                                       ProductFillBundle,
                                       ProductOrderElement,
                                       ProductOrderBundle)


class HuobiPrivateParser(object):
  @staticmethod
  def parse_balance(update_msg):
    assert update_msg['status'] == 'ok'
    currency_balances = {}
    for balance in update_msg['data']['list']:
      native_currency = balance['currency']
      currency = HuobiCurrency.FromStrNativeCurrencyNothrow(native_currency)
      if currency is None:
        continue
      if native_currency not in currency_balances:
        currency_balances[native_currency] = CurrencyBalance(
            currency=currency.currency, currency_native=currency.native_currency)
      currency_balance = currency_balances[native_currency]
      if balance['type'] == 'trade':
        currency_balance.available += float(balance['balance'])
      elif balance['type'] == 'frozen':
        currency_balance.hold += float(balance['balance'])
      elif balance['type'] in ('loan', 'interest'):
        currency_balance.available += float(balance['balance'])
      else:
        raise ValueError('Unexpected balance type!')

    for balance in currency_balances.values():
      balance.total = balance.available + balance.hold

    return AccountBalance(exchange='Huobi', each_balance=currency_balances.values())

  @staticmethod
  def parse_fill(update_msg, product):
    currency_fill_list = []
    for update_elem in update_msg['data']:
      if update_elem['symbol'] != product.native_symbol:
        continue
      if 'buy' in update_elem['type']:
        side = coin_order.BUY_ORDER
      elif 'sell' in update_elem['type']:
        side = coin_order.SELL_ORDER
      else:
        raise ValueError('Unknown order side! %s' % update_elem['type'])
      currency_fill = ProductFillElement(symbol=product.symbol,
                                         price=float(update_elem['price']),
                                         qty=float(update_elem['filled-amount']),
                                         side=side,
                                         fee=float(update_elem['filled-fees']),
                                         order_id=str(update_elem['order-id']),
                                         fill_timestamp=int(
                                             float(update_elem['created-at']) * 1e+6))
      currency_fill_list.append(currency_fill)
    return ProductFillBundle(exchange='Huobi', symbol=product.symbol, each_fill=currency_fill_list)

  @staticmethod
  def parse_orders(update_msg, product):
    currency_order_list = []
    for update_elem in update_msg['data']:
      if update_elem['symbol'] != product.native_symbol:
        continue
      if update_elem['type'] in ('sell-limit', 'sell-limit-maker'):
        side = coin_order.SELL_ORDER
      elif update_elem['type'] in ('buy-limit', 'buy-limit-maker'):
        side = coin_order.BUY_ORDER
      elif update_elem['type'] in ['sell-market', 'buy-market']:
        raise ValueError('We do not use %s order type.' % update_elem['type'])
      else:
        raise ValueError('Unknown order type %s ' % update_elem['type'])

      currency_order = ProductOrderElement(symbol=product.symbol,
                                           side=side,
                                           price=float(update_elem['price']),
                                           qty=float(update_elem['amount']),
                                           order_id=str(update_elem['id']),
                                           created_timestamp=int(
                                               float(update_elem['created-at']) * 1e+6))
      currency_order_list.append(currency_order)
    return ProductOrderBundle(exchange='Huobi',
                              symbol=product.symbol,
                              each_order=currency_order_list)

  @staticmethod
  def parse_account_id(update_msg):
    assert update_msg['status'] == "ok", update_msg
    res = {
        HuobiAccountType.SPOT: None,
        HuobiAccountType.MARGIN: [],
        HuobiAccountType.FUTURES: None,
        HuobiAccountType.POINT: None,
    }
    for account_info in update_msg['data']:
      huobi_account_type = HuobiAccountType(account_info['type'])
      if huobi_account_type == HuobiAccountType.MARGIN:
        res[huobi_account_type].append(account_info['id'])
      else:
        res[huobi_account_type] = account_info['id']
    return res

  @staticmethod
  def parse_transfers(update_msg, currency):
    """
    List of possible withdraw state

    State	Description
    submitted	Withdraw request submitted successfully
    reexamine	Under examination for withdraw validation
    canceled	Withdraw canceled by user
    pass	Withdraw validation passed
    reject	Withdraw validation rejected
    pre-transfer	Withdraw is about to be released
    wallet-transfer	On-chain transfer initiated
    wallet-reject	Transfer rejected by chain
    confirmed	On-chain transfer completed with one confirmation
    confirm-error	On-chain transfer faied to get confirmation
    repealed	Withdraw terminated by system
    """
    """
    List of possible deposit state

    State	Description
    unknown	On-chain transfer has not been received
    confirming	On-chain transfer waits for first confirmation
    confirmed	On-chain transfer confirmed for at least one block
    safe	Multiple on-chain confirmation happened
    orphan	Confirmed but currently in an orphan branch
    """
    status_map = {
        'unknown': 'TRANSFER_ACCEPTED',
        'confirming': 'TRANSFER_ACCEPTED',
        'confirmed': 'TRANSFER_CONFIRMED',
        'safe': 'TRANSFER_CONFIRMED',
        'orphan': 'TRANSFER_ACCEPTED',
        'submitted': 'TRANSFER_ACCEPTED',
        'reexamine': 'TRANSFER_ACCEPTED',
        'canceled': 'CANCEL_CONFIRMED',
        'pass': 'TRANSFER_CONFIRMED',
        'reject': 'TRANSFER_REJECTED',
        'pre-transfer': 'TRANSFER_ACCEPTED',
        'wallet-transfer': 'TRANSFER_CONFIRMED',
        'wallet-reject': 'TRANSFER_REJECTED',
        'confirmed': 'TRANSFER_CONFIRMED',
        'confirm-error': 'TRANSFER_REJECTED',
        'repealed': 'CANCEL_CONFIRMED'
    }
    transfer_list = []
    for msg in update_msg:
      assert msg['status'] == "ok", msg
      for transfer_elem in msg['data']:
        transfer = CurrencyTransferElement(currency=transfer_elem['currency'],
                                           type=transfer_elem['type'].upper(),
                                           to_address=transfer_elem['address'],
                                           qty=transfer_elem['amount'],
                                           fee=transfer_elem['fee'],
                                           external_transfer_id=transfer_elem['tx-hash'],
                                           status=status_map[transfer_elem['state']],
                                           create_timestamp=int(transfer_elem['created-at'] * 1e+6))
        transfer_list.append(transfer)
    return CurrencyTransferBundle(exchange='Huobi',
                                  market_type='Spot',
                                  currency=currency.currency,
                                  native_currency=currency.native_currency,
                                  each_transfer=transfer_list)


class HuobiPrivateClient(privcb.PrivateClientBase):
  ProductType = HuobiProduct
  _native_account_type = HuobiAccountType.SPOT

  def __init__(self, key_file):
    privcb.PrivateClientBase.__init__(self)
    self.nprivc = nprivc.HuobiNativePrivateClient(key_file=key_file)

  def query_native_account_id(self):
    update = self.nprivc.query_accounts()
    update.msg = HuobiPrivateParser.parse_account_id(update.msg)
    return update

  def query_account_balance_impl(self):
    account_update = self.query_native_account_id()
    native_account_id = account_update.msg[self._native_account_type]
    if native_account_id is None:
      raise RuntimeError('Cannot found account information for account %s'
                         % self._native_account_type)

    update = self.nprivc.query_account_balance(native_account_id)
    update.msg = HuobiPrivateParser.parse_balance(update.msg)

    # Compute the |total| field for each balance item.
    for balance in update.msg.each_balance:
      balance.total = balance.available + balance.hold
    return update

  def query_fills_impl(self, product):
    update = self.nprivc.query_order_match_results(product.native_symbol)
    update.msg = HuobiPrivateParser.parse_fill(update.msg, product)
    return update

  def query_list_orders_impl(self, product):
    update = self.nprivc.query_open_orders()
    update.msg = HuobiPrivateParser.parse_orders(update.msg, product)
    return update

  def query_account_transfer_impl(self, currency):
    currency = HuobiCurrency.FromStr(currency)
    update_msg = []
    transfer_types = ('deposit', 'withdraw')
    for type in transfer_types:
      update = self.nprivc.query_deposit_withdraw(currency.native_currency, type)
      update_msg.append(update.msg)
    update.msg = HuobiPrivateParser.parse_transfers(update_msg, currency)
    return update

  def query_submit_impl(self, product, side, price, qty):
    raise NotImplementedError()

  def query_cancel_product_impl(self, product):
    raise NotImplementedError()


class HuobiMarginPrivateClient(HuobiPrivateClient):
  ProductType = HuobiProduct
  _native_account_type = HuobiAccountType.MARGIN

  def query_account_balance_impl(self):
    account_update = self.query_native_account_id()
    native_account_ids = account_update.msg[self._native_account_type]
    if not native_account_ids:
      raise RuntimeError('Cannot found account information for account %s'
                         % self._native_account_type)
    total_balance = []
    for native_account_id in native_account_ids:
      update = self.nprivc.query_account_balance(native_account_id)
      total_balance.extend(update.msg['data']['list'])
    update.msg['data']['list'] = total_balance
    update.msg = HuobiPrivateParser.parse_balance(update.msg)
    return update


if __name__ == "__main__":
  client = HuobiPrivateClient(
      key_file='../../coin_key/view/prestoinv0/Huobi/view_key.json')
  print(client.query_account_transfer_impl('ETH').msg)
